from calendar import monthrange
from datetime import datetime, timedelta


def get_month_start_end(year_month_str, input_format="%Y%m", output_format="%Y%m%d") -> tuple[str]:
    """
    根据输入的年月字符串，输出指定格式的月初和月末日期

    Args:
        year_month_str (str): 年月字符串，如 "202305"
        input_format (str): 输入字符串的格式，默认为 "%Y%m"
        output_format (str): 输出日期的格式，默认为 "%Y%m%d"

    Returns:
        tuple: (月初日期, 月末日期) 格式化的字符串

    Examples:
        >>> get_month_start_end("2023-05")
        ('20230501', '20230531')

        >>> get_month_start_end("2023/05", "%Y/%m", "%Y年%m月%d日")
        ('2023年05月01日', '2023年05月31日')
    """
    # 解析输入的年月字符串
    date_obj = datetime.strptime(year_month_str, input_format)

    # 获取月初日期 (当月第一天)
    start_date = date_obj.replace(day=1)

    # 获取月末日期 (使用calendar.monthrange获取当月天数)
    _, last_day = monthrange(date_obj.year, date_obj.month)
    end_date = date_obj.replace(day=last_day)

    # 格式化输出
    start_date_str = start_date.strftime(output_format)
    end_date_str = end_date.strftime(output_format)

    return start_date_str, end_date_str


def get_date_range(start_date: str | datetime, end_date: str | datetime):
    """
    获取指定日期范围内的所有日期
    :param start_date: 开始日期，格式为 'yyyy-mm-dd'
    :param end_date: 结束日期，格式为 'yyyy-mm-dd'
    :return: 日期列表
    """
    # 解析日期字符串为 datetime 对象
    start_datetime = datetime.strptime(start_date, "%Y-%m-%d") if isinstance(start_date, str) else start_date
    end_datetime = datetime.strptime(end_date, "%Y-%m-%d") if isinstance(end_date, str) else end_date
    # 生成日期范围
    date_range = [start_datetime + timedelta(days=x) for x in range((end_datetime - start_datetime).days + 1)]
    # 格式化日期为 'yyyy-mm-dd' 格式
    formatted_dates = [date.strftime("%Y-%m-%d") for date in date_range]
    return formatted_dates
